///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/command.cpp
///
/// \brief   implementation file for fqmc related commands
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<string>

#include<formic/command/command.h>
#include<formic/fqmc/command.h>
#include<formic/fqmc/energy.h>
#include<formic/fqmc/optimizer.h>
#include<formic/all_data.h>
#include<formic/archive.h>
#include<formic/exception.h>
#include<formic/numeric/numeric.h>
#include<formic/mpi/interface.h>
#include<formic/fqmc/no_transform.h>
#include<formic/fqmc/orb_opt.h>
#include<formic/fqmc/onepdm.h>
#include<formic/fqmc/jastrow_utility.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function used to compute agp pairing matrix for the "AGP_Matrix_From_RHF" command
///
/// \param[in]     userinp      user input options
/// \param[in,out] wfn          the wavefunction
/// \param[in]     agp_name     name for the new agp pairing matrix object
/// \param[in]     rhf_name     name of the rhf orbital coefficient object to use
/// \param[in]     fixed        whether the pairing matrix will be fixed or variable
///
///////////////////////////////////////////////////////////////////////////////////////////////////

template <class S> static void formic_add_agp_matrix_from_rhf_detail(const formic::InputBase & userinp,
                                                                     formic::Wavefunction<S> & wfn,
                                                                     const std::string & agp_name,
                                                                     const std::string & rhf_name,
                                                                     const bool fixed)
{

  // get number of atomic orbitals
  const int nao = userinp.get<int>("nao");

  // get number of alpha particles
  const int na = userinp.get<int>("nalpha");
  if (na != userinp.get<int>("nbeta"))
    throw formic::Exception("command \"AGP_Matrix_From_RHF\" requires equal number of alpha and beta electrons");

  // get a pointer to the pairing matrix data object, creating and adding it to the wave funciton if necessary
  boost::shared_ptr<formic::WfnData<S> > pm;
  if ( wfn.has_wf(agp_name) ) {
    pm = wfn.get_wf(agp_name);
    if ( pm->data_vec().size() != nao*nao )
      throw formic::Exception("existing pairing matrix data object \"%s\" has size %i but should have size %i in formic_add_agp_matrix_from_rhf_detail") % agp_name % pm->data_vec().size() % (nao*nao);
  } else {
    boost::shared_ptr<formic::WfnData<S> > temp( new formic::WfnData<S>(agp_name, fixed, nao*nao, formic::zero(S())) );
    pm = temp;
    wfn.add_wfn_data(pm);
  }

  // get a reference to the rhf orbital coefficients
  const std::vector<S> & rvec = wfn.get_wf(rhf_name)->data_vec();

  // get a reference to the pairing matrix
  std::vector<S> & avec = pm->data_vec();

  // compute the pairing matrix
  std::fill(avec.begin(), avec.end(), formic::zero(S()));
  for (int p = 0; p < nao; p++)
  for (int q = 0; q < nao; q++)
  for (int i = 0; i < na; i++)
    avec.at(nao*p+q) += rvec.at(nao*i+p) * rvec.at(nao*i+q);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "AGP_Matrix_From_RHF" command, which creates a new data object
///          in the wavefunction containing the AGP pairing matrix converted from RHF orbital
///          coefficients.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_AGP_Matrix_From_RHF_command_function(formic::Archive & arch, formic::AllData & data) {

  // read in the name to use for the new agp pairing matrix data object
  std::string agp_name;
  formic::mpi::read_and_bcast(arch, agp_name, "Failed to read agp_name in formic_AGP_Matrix_From_RHF_command_function.");

  // read in the name of the object that holds the RHF orbital coefficients
  std::string rhf_name;
  formic::mpi::read_and_bcast(arch, rhf_name, "Failed to read rhf_name in formic_AGP_Matrix_From_RHF_command_function.");

  // read in whether the new data object will be fixed or variable
  bool fixed;
  formic::mpi::read_and_bcast(arch, fixed, "Failed to read fixed in formic_AGP_Matrix_From_RHF_command_function.");

  // create the data object
  if (data.userinp.get<std::string>("scalar_type") == "real")
    formic_add_agp_matrix_from_rhf_detail(data.userinp, data.wd, agp_name, rhf_name, fixed);
  else if (data.userinp.get<std::string>("scalar_type") == "complex")
    formic_add_agp_matrix_from_rhf_detail(data.userinp, data.wc, agp_name, rhf_name, fixed);
  else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"AGP_Matrix_From_RHF\"") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Implementation of the work done by the "AGP_Add_Double_Occ_Orb" command, in which
///          a singlet geminal corresponding to a doubly occupied orbital is added to the AGP
///          pairing matrix.
///
/// \param[in]     userinp      user input options
/// \param[in,out] wfn          the wavefunction
/// \param[in]     arch         archive to read command arguments from
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> static void formic_AGP_Add_Double_Occ_Orb_detail(const formic::InputBase & userinp,
                                                                    formic::Wavefunction<S> & wfn,
                                                                    formic::Archive & arch)
{

  // read in the name to use for the new agp pairing matrix data object
  std::string pm_name;
  formic::mpi::read_and_bcast(arch, pm_name, "Failed to read pm_name in formic_AGP_Add_Double_Occ_Orb_detail.");

  // read in the name to use for the new agp pairing matrix data object
  double weight;
  formic::mpi::read_and_bcast(arch, weight, "Failed to read weight in formic_AGP_Add_Double_Occ_Orb_detail.");

  // read in the name to use for the new agp pairing matrix data object
  formic::Vector<S,int> coeffs;
  formic::mpi::read_and_bcast(arch, coeffs, "Failed to read coeffs in formic_AGP_Add_Double_Occ_Orb_detail.");

  // get number of atomic orbitals
  const int nao = userinp.get<int>("nao");

  // check length of orbital coefficient vector
  if ( nao != coeffs.size() )
    throw formic::Exception("length of coefficient vector (%i) does not match number of basis orbitals (%i) in formic_AGP_Add_Double_Occ_Orb_detail") % coeffs.size() % nao;

  // get a reference to the pairing matrix
  std::vector<S> & pm = wfn.get_wf(pm_name)->data_vec();

  // ensure the pairing matrix is the correct size
  if ( nao*nao != pm.size() )
    throw formic::Exception("pairing matrix \"%s\" is the wrong size in formic_AGP_Add_Double_Occ_Orb_detail") % pm_name;

  // add the doubly occupied orbital geminal to the pairing matrix
  for (int p = 0; p < nao; p++)
  for (int q = 0; q < nao; q++)
    pm.at(nao*p+q) += weight * coeffs.at(p) * coeffs.at(q);

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "AGP_Add_Double_Occ_Orb" command, which adds a singlet
///          geminal corresponding to a doubly occupied orbital to the AGP pairing matrix.
///
///          Example usage:  AGP_Add_Double_Occ_Orb  pm  0.9  [ 0.0 0.7 0.7 0.0 0.0 ]
///
///          1st argument is the name of the agp pairing matrix data object
///          2nd argument is the (real) weight to multiply the added geminal by
///          3nd argument is a vector of orbital coefficients for the doubly occupied orbital
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_AGP_Add_Double_Occ_Orb_command_function(formic::Archive & arch, formic::AllData & data) {

  // call the appropriate template for either real or complex numbers
  if (data.userinp.get<std::string>("scalar_type") == "real")
    formic_AGP_Add_Double_Occ_Orb_detail(data.userinp, data.wd, arch);
  else if (data.userinp.get<std::string>("scalar_type") == "complex")
    formic_AGP_Add_Double_Occ_Orb_detail(data.userinp, data.wc, arch);
  else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"AGP_Add_Double_Occ_Orb\"") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "MCEnergy" command, which computes the energy of the named
///          coefficient object by MC sampling.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_MCEnergy_command_function(formic::Archive & arch, formic::AllData & data) {

  // read name of coefficient object whose energy we will compute
  std::string co_name;
  formic::mpi::read_and_bcast(arch, co_name, "Failed to read co_name in formic_MCEnergy_command_function.");

  // compute the energy
  if (data.userinp.get<std::string>("scalar_type") == "real")
    double energy = formic::fqmc::energy(data.userinp, data.wd.get_co(co_name));
  else if (data.userinp.get<std::string>("scalar_type") == "complex")
    double energy = formic::fqmc::energy(data.userinp, data.wc.get_co(co_name));
  else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"MCEnergy\"") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Optimize" command, which minimizes the energy of the
///          named coefficient object.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////

static void formic_Optimize_command_function(formic::Archive & arch, formic::AllData & data) {

  // read name of coefficient object to optimize
  std::string co_name;
  formic::mpi::read_and_bcast(arch, co_name, "Failed to read co_name in formic_Optimize_command_function.");

  // execute the optimization
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic::fqmc::optimize(data.userinp, data.wd.get_co(co_name));
    if (formic::mpi::rank() == 0)
      data.wd.dump_data_to_bin_arch("wfn_data.bin");
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic::fqmc::optimize(data.userinp, data.wc.get_co(co_name));
    if (formic::mpi::rank() == 0)
      data.wc.dump_data_to_bin_arch("wfn_data.bin");
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"Optimize\"") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   expands the orbital basis underneath the jastrow rgp data
///
/// \param[in]       userinp  the user's input options
/// \param[in,out]   wfn      the wavefunction object holding the matrices
/// \param[in,out]   arch     archive to read command arguments from
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S>
static void formic_BasisUpgradeJRGP_command_function_detail(const formic::InputBase & userinp,
                                                            formic::Wavefunction<S> & wfn,
                                                            formic::Archive & arch) {

  // check that simple jastrows are being used
  if ( ! userinp.get<bool>("simple_jastrow")  )
    throw formic::Exception("formic_BasisUpgradeJRGP_command_function_detail currently only implemented for simple jastrows");

  // get whether or not the jastrows are restricted
  const bool rest_jast = userinp.get<bool>("restrict_jastrow");

  // read in the particle number projection to use with the two-orbital mapping
  double proj_mag;
  formic::mpi::read_and_bcast(arch,  proj_mag, "Failed to read particle number projection magnitude in formic_BasisUpgradeJRGP_command_function_detail.");

  // read in the names of the small-basis data varaibles
  std::string spm_name, sjaa_name, sjbb_name, sjab_name;
  formic::mpi::read_and_bcast(arch,  spm_name, "Failed to read small pairing matrix data object name in formic_BasisUpgradeJRGP_command_function_detail.");
  formic::mpi::read_and_bcast(arch, sjaa_name, "Failed to read small aa jastrow data object name in formic_BasisUpgradeJRGP_command_function_detail.");
  if ( !rest_jast )
    formic::mpi::read_and_bcast(arch, sjbb_name, "Failed to read small bb jastrow data object name in formic_BasisUpgradeJRGP_command_function_detail.");
  formic::mpi::read_and_bcast(arch, sjab_name, "Failed to read small ab jastrow data object name in formic_BasisUpgradeJRGP_command_function_detail.");

  // read in the names of the large-basis data varaibles
  std::string lpm_name, ljaa_name, ljbb_name, ljab_name;
  formic::mpi::read_and_bcast(arch,  lpm_name, "Failed to read large pairing matrix data object name in formic_BasisUpgradeJRGP_command_function_detail.");
  formic::mpi::read_and_bcast(arch, ljaa_name, "Failed to read large aa jastrow data object name in formic_BasisUpgradeJRGP_command_function_detail.");
  if ( !rest_jast )
    formic::mpi::read_and_bcast(arch, ljbb_name, "Failed to read large bb jastrow data object name in formic_BasisUpgradeJRGP_command_function_detail.");
  formic::mpi::read_and_bcast(arch, ljab_name, "Failed to read large ab jastrow data object name in formic_BasisUpgradeJRGP_command_function_detail.");

  // read in the orbital mapping
  std::vector<int> mapping;
  formic::mpi::read_and_bcast(arch, mapping, "Failed to read orbital mapping in formic_BasisUpgradeJRGP_command_function_detail.");

  // get a reference to the data vectors
  std::vector<S> &  spm_vec = wfn.get_wf( spm_name)->data_vec();
  std::vector<S> & sjaa_vec = wfn.get_wf(sjaa_name)->data_vec();
  std::vector<S> & sjbb_vec = ( rest_jast ? sjaa_vec : wfn.get_wf(sjbb_name)->data_vec() );
  std::vector<S> & sjab_vec = wfn.get_wf(sjab_name)->data_vec();
  std::vector<S> &  lpm_vec = wfn.get_wf( lpm_name)->data_vec();
  std::vector<S> & ljaa_vec = wfn.get_wf(ljaa_name)->data_vec();
  std::vector<S> & ljbb_vec = ( rest_jast ? ljaa_vec : wfn.get_wf(ljbb_name)->data_vec() );
  std::vector<S> & ljab_vec = wfn.get_wf(ljab_name)->data_vec();

  // get number of orbitals in small basis set
  int ns = 0;
  for ( ; ns < spm_vec.size(); ns++) {
    if ( ns * ns == spm_vec.size() )
      break;
    if ( ns * ns > spm_vec.size() )
      throw formic::Exception("The data in \"%s\" is not a square matrix in formic_BasisUpgradeJRGP_command_function_detail.") % spm_name;
  }

  // get number of orbitals in large basis set
  int nl = 0;
  for ( ; nl < lpm_vec.size(); nl++) {
    if ( nl * nl == lpm_vec.size() )
      break;
    if ( nl * nl > lpm_vec.size() )
      throw formic::Exception("The data in \"%s\" is not a square matrix in formic_BasisUpgradeJRGP_command_function_detail.") % lpm_name;
  }

  // translate the mapping
  std::vector<int> map_num(ns, -1); // whether it is a one- or two-orbital mapping
  std::vector<int> map_1st(ns, -1); // first orbital to map to
  std::vector<int> map_2nd(ns, -1); // second orbital to map to
  for (std::vector<int>::const_iterator it = mapping.begin(); it != mapping.end(); ) {

    // get whether this is a single or double mapping
    const int num = -(*it++);
    if ( num != 1 && num != 2 )
      throw formic::Exception("mapping type was \"%i\" when it should be either \"-1\" or \"-2\" in formic_BasisUpgradeJRGP_command_function_detail") % num;

    // get the small-basis orbital number
    if ( it == mapping.end() )
      throw formic::Exception("unexpectedly reached the end of the mapping vector in formic_BasisUpgradeJRGP_command_function_detail");
    const int so = *it++;
    if ( so < 0 || so >= ns )
      throw formic::Exception("out-of-bounds small-basis orbital number \"%i\" in formic_BasisUpgradeJRGP_command_function_detail") % so;
    map_num.at(so) = num;

    // get the first large-basis orbital number
    if ( it == mapping.end() )
      throw formic::Exception("unexpectedly reached the end of the mapping vector in formic_BasisUpgradeJRGP_command_function_detail");
    map_1st.at(so) = *it++;
    if ( map_1st.at(so) < 0 || map_1st.at(so) >= nl )
      throw formic::Exception("out-of-bounds large-basis orbital number \"%i\" in formic_BasisUpgradeJRGP_command_function_detail") % map_1st.at(so);

    // if this is a double mapping, get the second large-basis orbital number
    if ( num == 2 ) {
      if ( it == mapping.end() )
        throw formic::Exception("unexpectedly reached the end of the mapping vector in formic_BasisUpgradeJRGP_command_function_detail");
      map_2nd.at(so) = *it++;
      if ( map_2nd.at(so) < 0 || map_2nd.at(so) >= nl )
        throw formic::Exception("out-of-bounds large-basis orbital number \"%i\" in formic_BasisUpgradeJRGP_command_function_detail") % map_2nd.at(so);
    }

  }

  // expand the pairing matrix
  std::fill(lpm_vec.begin(), lpm_vec.end(), formic::zero(S()));
  for (int i = 0; i < ns; i++)
  for (int j = 0; j < ns; j++) {

    // single-single mapping
    if ( map_num.at(i) == 1 && map_num.at(j) == 1 ) {

      lpm_vec.at( nl * map_1st.at(i) + map_1st.at(j) ) += spm_vec.at( ns * i + j );

    // single-double mapping
    } else if ( map_num.at(i) == 1 && map_num.at(j) == 2 ) {

      lpm_vec.at( nl * map_1st.at(i) + map_1st.at(j) ) += spm_vec.at( ns * i + j ) / std::sqrt(2.0);
      lpm_vec.at( nl * map_1st.at(i) + map_2nd.at(j) ) += spm_vec.at( ns * i + j ) / std::sqrt(2.0);

    // double-single mapping
    } else if ( map_num.at(i) == 2 && map_num.at(j) == 1 ) {

      lpm_vec.at( nl * map_1st.at(i) + map_1st.at(j) ) += spm_vec.at( ns * i + j ) / std::sqrt(2.0);
      lpm_vec.at( nl * map_2nd.at(i) + map_1st.at(j) ) += spm_vec.at( ns * i + j ) / std::sqrt(2.0);

    // double-double mapping
    } else if ( map_num.at(i) == 2 && map_num.at(j) == 2 ) {

      lpm_vec.at( nl * map_1st.at(i) + map_1st.at(j) ) += spm_vec.at( ns * i + j ) / 2.0;
      lpm_vec.at( nl * map_2nd.at(i) + map_1st.at(j) ) += spm_vec.at( ns * i + j ) / 2.0;
      lpm_vec.at( nl * map_1st.at(i) + map_2nd.at(j) ) += spm_vec.at( ns * i + j ) / 2.0;
      lpm_vec.at( nl * map_2nd.at(i) + map_2nd.at(j) ) += spm_vec.at( ns * i + j ) / 2.0;

    // error
    } else {
      throw formic::Exception("unexpected pair of map_num values (%i,%i) in formic_BasisUpgradeJRGP_command_function_detail") % map_num.at(i) % map_num.at(j);
    }

  }

  // define some local functions to help with jastrow expansion
  struct local_funcs {

    // define local function to handle same-spin jastrow expansion
    static void expand_same_spin_aa(const int ns,
                                    const int nl,
                                    const std::vector<int> & map_num,
                                    const std::vector<int> & map_1st,
                                    const std::vector<int> & map_2nd,
                                    const double proj_mag,
                                    const std::vector<S> & sjaa_vec,
                                    std::vector<S> & ljaa_vec) {

      // zero the large jastrow factor
      std::fill(ljaa_vec.begin(), ljaa_vec.end(), formic::zero(S()));

      // loop over first small basis index
      for (int i = 0; i < ns; i++) {

        // partial number projection
        if ( map_num.at(i) == 2 ) {
          const int p = std::min(map_1st.at(i), map_2nd.at(i));
          const int q = std::max(map_1st.at(i), map_2nd.at(i));
          ljaa_vec.at( formic::cmpd_pair_index_incl(p, q, nl) ) -= 2.0 * proj_mag * formic::unity(S());
        }

        // loop over second small basis index
        for (int j = i; j < ns; j++) {

          // get small basis jastrow value for orbitals i and j
          const S x = sjaa_vec.at( formic::cmpd_pair_index_incl(i, j, ns) );

          // single-single mapping
          if ( map_num.at(i) == 1 && map_num.at(j) == 1 ) {

            ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_1st.at(i), map_1st.at(j)), std::max(map_1st.at(i), map_1st.at(j)), nl) ) += x;

          // single-double mapping
          } else if ( map_num.at(i) == 1 && map_num.at(j) == 2 ) {

            ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_1st.at(i), map_1st.at(j)), std::max(map_1st.at(i), map_1st.at(j)), nl) ) += x;
            ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_1st.at(i), map_2nd.at(j)), std::max(map_1st.at(i), map_2nd.at(j)), nl) ) += x;

          // double-single mapping
          } else if ( map_num.at(i) == 2 && map_num.at(j) == 1 ) {

            ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_1st.at(i), map_1st.at(j)), std::max(map_1st.at(i), map_1st.at(j)), nl) ) += x;
            ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_2nd.at(i), map_1st.at(j)), std::max(map_2nd.at(i), map_1st.at(j)), nl) ) += x;

          // double-double mapping
          } else if ( map_num.at(i) == 2 && map_num.at(j) == 2 ) {

            ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_1st.at(i), map_1st.at(j)), std::max(map_1st.at(i), map_1st.at(j)), nl) ) += x;
            ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_2nd.at(i), map_2nd.at(j)), std::max(map_2nd.at(i), map_2nd.at(j)), nl) ) += x;
            if ( i != j ) {
              ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_1st.at(i), map_2nd.at(j)), std::max(map_1st.at(i), map_2nd.at(j)), nl) ) += x;
              ljaa_vec.at( formic::cmpd_pair_index_incl( std::min(map_2nd.at(i), map_1st.at(j)), std::max(map_2nd.at(i), map_1st.at(j)), nl) ) += x;
            }

          // error
          } else {
            throw formic::Exception("unexpected pair of map_num values (%i,%i) in same-spin jastrow portion of formic_BasisUpgradeJRGP_command_function_detail") % map_num.at(i) % map_num.at(j);
          }

        }

      }

    }

  };

  // expand same-spin jastrows
  local_funcs::expand_same_spin_aa(ns, nl, map_num, map_1st, map_2nd, proj_mag, sjaa_vec, ljaa_vec);
  if ( !rest_jast )
    local_funcs::expand_same_spin_aa(ns, nl, map_num, map_1st, map_2nd, proj_mag, sjbb_vec, ljbb_vec);

  // expand opposite-spin jastrows
  std::fill(ljab_vec.begin(), ljab_vec.end(), formic::zero(S()));
  for (int i = 0; i < ns; i++)
  for (int j = ( rest_jast ? i : 0 ); j < ns; j++) {

    const S x = sjab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, i, j, ns) );

    // single-single mapping
    if ( map_num.at(i) == 1 && map_num.at(j) == 1 ) {

      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_1st.at(i), map_1st.at(j), nl) ) += x;

    // single-double mapping
    } else if ( map_num.at(i) == 1 && map_num.at(j) == 2 ) {

      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_1st.at(i), map_1st.at(j), nl) ) += x;
      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_1st.at(i), map_2nd.at(j), nl) ) += x;

    // double-single mapping
    } else if ( map_num.at(i) == 2 && map_num.at(j) == 1 ) {

      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_1st.at(i), map_1st.at(j), nl) ) += x;
      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_2nd.at(i), map_1st.at(j), nl) ) += x;

    // double-double mapping
    } else if ( map_num.at(i) == 2 && map_num.at(j) == 2 ) {

      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_1st.at(i), map_1st.at(j), nl) ) += x;
      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_2nd.at(i), map_2nd.at(j), nl) ) += x;
      ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_1st.at(i), map_2nd.at(j), nl) ) += x;
      if ( i != j || !rest_jast )
        ljab_vec.at( formic::fqmc::jastrow_ab_cmpd_ind(rest_jast, map_2nd.at(i), map_1st.at(j), nl) ) += x;

    // error
    } else {
      throw formic::Exception("unexpected pair of map_num values (%i,%i) in ab jastrow part of formic_BasisUpgradeJRGP_command_function_detail") % map_num.at(i) % map_num.at(j);
    }

  }

  // broadcast the large basis data vectors to all processes
  formic::mpi::bcast(lpm_vec);
  formic::mpi::bcast(ljaa_vec);
  if ( !rest_jast )
    formic::mpi::bcast(ljbb_vec);
  formic::mpi::bcast(ljab_vec);

  // empty any exact sample data in the wave function that may now have been invalidated
  wfn.clear_exact_sample_vecs();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "BasisUpgradeJRGP" command, which expands the orbital basis
///          for a jastrow agp.
///          Each original basis orbital is either mapped to one or two new basis orbitals.
///          In the case of mapping to two new basis orbitals, it is assumed that the old orbital
///          is equal to the normalized sum of the two new orbitals.
///
///          The first argument is the number projection penalty to use for two-orbital mappings.
///
///          The next arguments are the names of the small-basis jastrow rgp data objects,
///          e.g.   spm sjaa sjab          (for   restricted jastrows)
///          or     spm sjaa sjbb sjab     (for unrestricted jastrows)
///
///          The next arguments are the names of the large-basis jastrow rgp data objects,
///          e.g.   lpm ljaa ljab          (for   restricted jastrows)
///          or     lpm ljaa ljbb ljab     (for unrestricted jastrows)
///
///          The last argument is a vector of integers, which specifies the mapping.
///          Note that the number of mapping orbitals are expressed with negative numbers.
///          e.g.
///          %  number of mapping orbitals      original orb      new orb 1      new orb 2
///          [
///                    -1                           0                 0
///                    -2                           1                 1              2
///                    -2                           2                 3              4
///                    -2                           3                 5              6
///          ]
///
/// \param[in,out]   arch     archive to read command arguments from
/// \param[in,out]   data     data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
static void formic_BasisUpgradeJRGP_command_function(formic::Archive & arch, formic::AllData & data) {

  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic_BasisUpgradeJRGP_command_function_detail(data.userinp, data.wd, arch);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic_BasisUpgradeJRGP_command_function_detail(data.userinp, data.wc, arch);
  } else 
    throw formic::Exception("Unknown scalar type \"%s\" when processing BasisUpgradeJRGP command.") % data.userinp.get<std::string>("scalar_type");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Add fqmc related commands to the formic command system
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void formic::init_fqmc_commands() {

  // add command to add an agp pairing matrix to the wavefunction data, derived from RHF coefficients
  formic::add_command("AGP_Matrix_From_RHF", &formic_AGP_Matrix_From_RHF_command_function);

  // add command to add a singlet pair geminal consisting of a doubly occupied orbital (e.g. a bond, antibond, etc.) to the agp pairing matrix
  formic::add_command("AGP_Add_Double_Occ_Orb", &formic_AGP_Add_Double_Occ_Orb_command_function);

  // add command to compute the energy of a coefficient value object
  formic::add_command("MCEnergy", &formic_MCEnergy_command_function);

  // add command to optimize a coefficient value object
  formic::add_command("Optimize", &formic_Optimize_command_function);

  // add command to do a transformation to select natural orbitals
  formic::add_command("Select_NO_Transform", &formic::fqmc::Select_NO_Transform_command_function);

  // add command to do an arbitrary orbital transformation
  formic::add_command("Select_Orb_Transform", &formic::fqmc::Select_Orb_Transform_command_function);

  // add command to optimize the one electron orbital basis
  formic::add_command("OptimizeOrbitals", &formic::fqmc::OptimizeOrbitals_command_function);

  // add command to print natural orbitals
  formic::add_command("PrintNO", &formic::fqmc::PrintNO_command_function);

  // add command to upgrade a JRGP wave function from a smaller to a larger basis
  formic::add_command("BasisUpgradeJRGP", &formic_BasisUpgradeJRGP_command_function);

}
